#! /usr/bin/env python
import os, sys
import rospy
from sensor_msgs.msg import Image
import tf2_ros
import numpy as np
from tf2_geometry_msgs import do_transform_pose
from visualization_msgs.msg import MarkerArray
from grasp_generation.srv import GraspGenerate

from moveit_commander import MoveGroupCommander
from sr_robot_commander.sr_hand_commander import SrHandCommander

HERE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(HERE_DIR, '../../utils/'))
import ros_np_multiarray as rosnm
from transform import R2q, q2R
from msg_converter import get_posestamped_msg

def open_shadow_hand(hand_commander):
    open_hand = {'rh_FFJ1': 0.0, 'rh_FFJ2': 0.0, 'rh_FFJ3': 0.0, 'rh_FFJ4': 0.0,
                'rh_MFJ1': 0.0, 'rh_MFJ2': 0.0, 'rh_MFJ3': 0.0, 'rh_MFJ4': 0.0,
                'rh_RFJ1': 0.0, 'rh_RFJ2': 0.0, 'rh_RFJ3': 0.0, 'rh_RFJ4': 0.0,
                'rh_LFJ1': 0.0, 'rh_LFJ2': 0.0, 'rh_LFJ3': 0.0, 'rh_LFJ4': 0.0, 'rh_LFJ5': 0.0,
                'rh_THJ1': 0.0, 'rh_THJ2': 0.0, 'rh_THJ3': 0.0, 'rh_THJ4': 0.0, 'rh_THJ5': 0.0,
                'rh_WRJ1': 0.0, 'rh_WRJ2': 0.0}
    # Move hand to open position
    joint_states = open_hand
    rospy.loginfo("Moving hand to open position")
    hand_commander.move_to_joint_value_target_unsafe(joint_states, 2.0, False)
    rospy.sleep(2)

def define1_shadow_hand(hand_commander):
    defined_pose = {'rh_FFJ1': 0, 'rh_FFJ2': 0.0, 'rh_FFJ3': 1.2, 'rh_FFJ4': 0.0,
                'rh_MFJ1': 0, 'rh_MFJ2': 0.0, 'rh_MFJ3': 1.2, 'rh_MFJ4': 0.0,
                'rh_RFJ1': 1.5707, 'rh_RFJ2': 1.5707, 'rh_RFJ3': 1.5707, 'rh_RFJ4': 0.0,
                'rh_LFJ1': 1.5707, 'rh_LFJ2': 1.5707, 'rh_LFJ3': 1.5707, 'rh_LFJ4': 0.0, 'rh_LFJ5': 0.0,
                'rh_THJ1': 0.0, 'rh_THJ2': 0.0, 'rh_THJ3': 0.0, 'rh_THJ4': 0.0, 'rh_THJ5': 0.0,
                'rh_WRJ1': 0.0, 'rh_WRJ2': 0.0}
    joint_states = defined_pose
    rospy.loginfo("Moving hand to define1 position")
    hand_commander.move_to_joint_value_target_unsafe(joint_states, 2.0, False)
    rospy.sleep(2)

def define2_shadow_hand(hand_commander):
    defined_pose = {'rh_FFJ1': 0.0, 'rh_FFJ2': 0.0, 'rh_FFJ3': 0.0, 'rh_FFJ4': 0.0,
                'rh_MFJ1': 0.0, 'rh_MFJ2': 0.0, 'rh_MFJ3': 0.0, 'rh_MFJ4': 0.0,
                'rh_RFJ1': 0.0, 'rh_RFJ2': 0.0, 'rh_RFJ3': 0.0, 'rh_RFJ4': 0.0,
                'rh_LFJ1': 0.0, 'rh_LFJ2': 0.0, 'rh_LFJ3': 0.0, 'rh_LFJ4': 0.0, 'rh_LFJ5': 0.0,
                'rh_THJ1': 0.0, 'rh_THJ2': 0.0, 'rh_THJ3': 0.0, 'rh_THJ4': 1.17, 'rh_THJ5': 0.0,
                'rh_WRJ1': 0.0, 'rh_WRJ2': 0.0}
    joint_states = defined_pose
    rospy.loginfo("Moving hand to define2 position")
    hand_commander.move_to_joint_value_target_unsafe(joint_states, 2.0, False)
    rospy.sleep(2)    

def define3_shadow_hand(hand_commander):
    defined_pose = { 'rh_FFJ1': 1.1,
                'rh_MFJ1': 1.1, 
                'rh_RFJ1': 1.1, 
                'rh_LFJ1': 1.1,
                'rh_THJ1': 0.7, 'rh_THJ2': 0.29, 'rh_THJ5': 0.35}
    joint_states = defined_pose
    rospy.loginfo("Moving hand to define3 position")
    hand_commander.move_to_joint_value_target_unsafe(joint_states, 2.0, True)
    rospy.sleep(2)


def close_shadow_hand(hand_commander):
    pack_hand_1 = {'rh_FFJ1': 1.5707, 'rh_FFJ2': 1.5707, 'rh_FFJ3': 1.5707, 'rh_FFJ4': 0.0,
                'rh_MFJ1': 1.5707, 'rh_MFJ2': 1.5707, 'rh_MFJ3': 1.5707, 'rh_MFJ4': 0.0,
                'rh_RFJ1': 1.5707, 'rh_RFJ2': 1.5707, 'rh_RFJ3': 1.5707, 'rh_RFJ4': 0.0,
                'rh_LFJ1': 1.5707, 'rh_LFJ2': 1.5707, 'rh_LFJ3': 1.5707, 'rh_LFJ4': 0.0, 'rh_LFJ5': 0.0}
    pack_hand_2 = {'rh_THJ4': 1.17}
    pack_hand_3 = {'rh_THJ1': 0.97, 'rh_THJ2': 0.29, 'rh_THJ5': 0.35}
    # open shadow hand
    open_shadow_hand(hand_commander)
    # Move hand to closed position
    joint_states = pack_hand_1
    rospy.loginfo("Moving hand to pack position")
    hand_commander.move_to_joint_value_target_unsafe(joint_states, 2.0, True)
    rospy.sleep(1)
    joint_states = pack_hand_2
    hand_commander.move_to_joint_value_target_unsafe(joint_states, 1.0, True)
    rospy.sleep(1)
    joint_states = pack_hand_3
    hand_commander.move_to_joint_value_target_unsafe(joint_states, 1.0, True)
    rospy.sleep(1)


def arm_plan_and_go(arm_commander, target_pose):
    arm_commander.set_pose_target(target_pose)
    result = arm_commander.plan()
    # rospy.loginfo('plan result : %s' % result)
    go_flag = raw_input('Execute ? y(es) / n(o) :')
    return go_flag == 'y'

if __name__ == '__main__':
    rospy.init_node('grasp_execute', anonymous=True)
    tfBuffer = tf2_ros.Buffer()
    tf_listenser = tf2_ros.TransformListener(tfBuffer)

    camera_frame = 'camera_color_optical_frame'

    hand = SrHandCommander(name='right_hand')
    arm = MoveGroupCommander('right_arm_and_wrist')
    arm.set_max_velocity_scaling_factor(0.05)
    arm.set_max_acceleration_scaling_factor(0.05)    

    # arm.set_goal_position_tolerance(0.02)
    # arm.set_goal_orientation_tolerance(0.1)
    
    keys = [           'rh_FFJ4','rh_FFJ3','rh_FFJ2','rh_FFJ1',
                       'rh_MFJ4','rh_MFJ3','rh_MFJ2','rh_MFJ1',
                       'rh_RFJ4','rh_RFJ3','rh_RFJ2','rh_RFJ1',
            'rh_LFJ5', 'rh_LFJ4','rh_LFJ3','rh_LFJ2','rh_LFJ1',
            'rh_THJ5', 'rh_THJ4','rh_THJ3','rh_THJ2','rh_THJ1']


    rospy.loginfo('go initial pose')
    initial_joint_value_ur10 = [1.7589348554611206, -0.8362162748919886, -1.9958394209491175, 
                                2.310129165649414, -1.3493455092059534, -1.4972642103778284, 
                                0.015987684898372913, 0.10608307679692915]
    arm.set_joint_value_target(initial_joint_value_ur10)
    arm.plan()
    if raw_input('execute? y(es) / n(o):') == 'y':
        arm.go()     

    close_shadow_hand(hand)

    grasp_client = rospy.ServiceProxy('/dexgrasp/grasp_server', GraspGenerate)
    grasp_marker_puber = rospy.Publisher('/dexgrasp/grasp', MarkerArray, queue_size=1)
    
    while not rospy.is_shutdown():
        img_raw = rospy.wait_for_message("/camera/color/image_raw", Image)
        depth_raw = rospy.wait_for_message("/camera/aligned_depth_to_color/image_raw", Image)    
        result = grasp_client(img_raw, depth_raw)
        if len(result.grasp_marker.markers) == 0:
            rospy.sleep(1.0)
            continue
        grasp_marker_puber.publish(result.grasp_marker)

        plan_flag = raw_input("A good grasp ? Go on now ? y(es) / n(o) :")
        if plan_flag == "y":
            palm_q = rosnm.to_numpy_f32(result.q)
            palm_t = rosnm.to_numpy_f32(result.t)
            obj_q = rosnm.to_numpy_f32(result.obj_q)
            obj_t = rosnm.to_numpy_f32(result.obj_t)
            palm_R = q2R(palm_q)
            obj_R = q2R(obj_q)
            grasp_pose = get_posestamped_msg(position=palm_t, orientation=palm_q, frame_name=camera_frame)
            pre_plam_t = palm_t + 0.15 * obj_R[:3, 1]  # grasp mug handle
            pre_grasp_pose = get_posestamped_msg(position=pre_plam_t, orientation=palm_q, frame_name=camera_frame)
            take_up_pose = get_posestamped_msg(position=palm_t, orientation=palm_q, frame_name=camera_frame)
            transform_ok = False
            while not transform_ok and not rospy.is_shutdown():
                try:
                    transform = tfBuffer.lookup_transform("ra_base_link", pre_grasp_pose.header.frame_id, rospy.Time(0))
                    pre_grasp_pose = do_transform_pose(pre_grasp_pose, transform)
                    transform = tfBuffer.lookup_transform("ra_base_link", grasp_pose.header.frame_id, rospy.Time(0))
                    grasp_pose = do_transform_pose(grasp_pose, transform)
                    transform = tfBuffer.lookup_transform("ra_base_link", take_up_pose.header.frame_id, rospy.Time(0))
                    take_up_pose = do_transform_pose(take_up_pose, transform)
                    transform_ok = True
                except tf2_ros.ExtrapolationException as e:
                    rospy.logwarn( "Exception on transforming point... trying again \n(" +str(e) + ")")
                    rospy.sleep(0.01)
                    pre_grasp_pose.header.stamp = tfBuffer.get_latest_common_time("ra_base_link", grasp_pose.header.frame_id)
                    grasp_pose.header.stamp = tfBuffer.get_latest_common_time("ra_base_link", grasp_pose.header.frame_id)
                    take_up_pose.header.stamp = tfBuffer.get_latest_common_time("ra_base_link", grasp_pose.header.frame_id)

            # pre_grasp_pose.pose.position.z += 0.15
            take_up_pose.pose.position.z += 0.15

            rospy.loginfo('go pre grasp')
            if arm_plan_and_go(arm, pre_grasp_pose):
                arm.go()
            else:
                continue

            rospy.loginfo('open hand')
            open_shadow_hand(hand)
            # nie_shadow_hand(hand)
            # define2_shadow_hand(hand)
            rospy.sleep(1.0)
            
            rospy.loginfo('grasp')
            if arm_plan_and_go(arm, grasp_pose):
                arm.go()

            rospy.loginfo('close hand')
            theta = rosnm.to_numpy_f32(result.theta)
            hand_target_pose = dict(zip(keys, theta))
            # There are some difference between openAI shadow hand and real shadow hand.
            hand_target_pose['rh_THJ1'] *= -1
            hand_target_pose['rh_THJ2'] *= -1
            hand_target_pose['rh_FFJ4'] *= -1
            hand_target_pose['rh_MFJ4'] *= -1
            hand.move_to_joint_value_target_unsafe(hand_target_pose, time=1.0, wait=True, angle_degrees=False)
            rospy.sleep(1.0)
            # define3_shadow_hand(hand)

            rospy.loginfo('take up')
            if arm_plan_and_go(arm, take_up_pose):
                arm.go()

            rospy.loginfo('put down')
            if arm_plan_and_go(arm, grasp_pose):
                arm.go()

            rospy.loginfo('open hand')
            open_shadow_hand(hand)

            rospy.loginfo('pre grasp hand')
            if arm_plan_and_go(arm, pre_grasp_pose):
                arm.go()

            rospy.loginfo('go back')
            arm.set_joint_value_target(initial_joint_value_ur10)
            arm.plan()
            if raw_input('execute? y(es) / n(o):') == 'y':
                arm.go()                 

            rospy.loginfo('close hand')
            close_shadow_hand(hand)